import asyncio
import websockets
from cozepy import Coze, TokenAuth, Message, COZE_CN_BASE_URL
import json
from config import CONFIG 

async def call_coze_workflow(workflow_id: str, input_text: str):
    """改进后的工作流调用函数"""
    try:
        coze = Coze(
            auth=TokenAuth(token=CONFIG["api_token"]),
            base_url=COZE_CN_BASE_URL
        )
        message = Message.build_user_question_text(input_text)
        
        loop = asyncio.get_running_loop()
        response = await loop.run_in_executor(
            None,
            lambda: coze.workflows.runs.create(
                workflow_id=workflow_id,
                parameters={'input': message.content}
            )
        )
        
        # 调试输出原始响应
        print("\n[原始响应] 工作流返回数据:")
        print(f"数据类型: {type(response.data)}")
        print(f"原始内容: {response.data}")

        # 新版响应处理逻辑
        processed_data = {}
        if isinstance(response.data, dict):
            processed_data = response.data
        elif isinstance(response.data, str):
            try:
                processed_data = json.loads(response.data)
            except json.JSONDecodeError:
                processed_data = {"output": response.data}  # 原始字符串包装成字典
        else:
            processed_data = {"raw_data": str(response.data)}

        return {
            "status": "success",
            "data": processed_data,
            "error": None
        }

    except Exception as e:
        print(f"API调用异常: {str(e)}")
        return {
            "status": "error",
            "data": None,
            "error": str(e)
        }

async def handle_server_messages(websocket):
    """处理来自服务器的消息"""
    try:
        while True:
            response = await websocket.recv()
            print("\n[服务器响应]:")
            print(json.dumps(json.loads(response), indent=2, ensure_ascii=False))
    except websockets.exceptions.ConnectionClosed:
        print("\n[信息] 服务器连接已关闭")
    except Exception as e:
        print(f"\n[错误] 接收消息时出错: {str(e)}")

async def main():
    async with websockets.connect(CONFIG["server_ws"]) as ws:
        # 启动接收消息的任务
        receive_task = asyncio.create_task(handle_server_messages(ws))
        
        try:
            while True:
                # 获取用户输入
                user_input = input("\n请输入学生信息（按q退出）: ").strip()
                
                if user_input.lower() == 'q':
                    print("\n[信息] 正在退出程序...")
                    break
                
                if not user_input:
                    continue
                
                # 调用第一个Bot
                print("\n[第一阶段] 正在处理学生需求...")
                first_result = await call_coze_workflow(CONFIG["bot1_workflow_id"], user_input)
                
                if first_result["status"] != "success":
                    print("❗ 第一阶段处理失败:", first_result["error"])
                    continue
                
                try:
                    output_content = first_result["data"]["output"]
                    if not output_content.strip():
                        raise ValueError("output字段内容为空")
                except (KeyError, ValueError) as e:
                    print(f"⚠ 数据异常：{str(e)}")
                    continue
                
                # 发送数据到服务器
                print("\n[第二阶段] 提交数据到服务器...")
                await ws.send(json.dumps({
                    "input": output_content
                }))
                print("已提交数据到服务器，等待响应...")
                
        except Exception as e:
            print(f"\n[错误] 发生异常: {str(e)}")
        finally:
            # 取消接收消息的任务
            receive_task.cancel()
            try:
                await receive_task
            except asyncio.CancelledError:
                pass

if __name__ == "__main__":
    try:
        asyncio.run(main())
    except KeyboardInterrupt:
        print("\n[信息] 程序被用户中断")
    except Exception as e:
        print(f"\n[错误] 程序异常退出: {str(e)}")